'''
	Author：Yassin
    用于交易回测，默认分红后再投资。
    以日/周/月/年为周期进行回测,返回周期内每一天的收益率。
'''

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

from trade_tools.Sel_from_ret import *
from trade_tools.Sel_from_report import *

def trade_period(Portfolio:list, Start:str, Hold_days : int = 0, Hold_period : str = '', Hold_method : str = 'Equal'):
    '''
        用于交易回测， 包括必选参数 Portfolio 和 Start， 以及可选参数 Hold_days 和 Hold_period
        Portfolio：list 所选组合的股票代码， int 形式
        Start：str  回测的开始日期，形式为 YYYYMMDD
        Hold_days : int 以日为单位回测的持股周期，需要大于等于 1，Start 日买入， 持有Hold_days个交易日，收盘价卖出
        Hold_period: str 以自然月"Month"/季度"Season"/年"Year"为单位持股周期，在周期的最后一天卖出，周期的第一天买入
        注意，季度持仓默认不跨年，如果需要跨年，可以自己用month连续做三次。
        Hold_method: str 默认为 Equal， 等金额购买，可选 Adapt，按照总市值比例购买
        只循环一个周期
        返回的 DataFrame 包括三列，Trddt，每日收益率 Daily_ret 以及累计收益率 Acc_ret
    '''

    ## 获取交易日历，找出时间范围内的第一个交易日

    calen = Get_trade_dates(Start, '20221231')
    Start_time = pd.to_datetime(Start, format = "%Y%m%d")

    calen = calen[calen["cal_date"] >= Start_time]
    calen = calen.sort_values("cal_date", ascending = True)
    calen = calen.reset_index(drop = True)
    Time_list = list(calen["cal_date"])
    Start = Time_list[0]

    ## 计算个股权重

    Weights = []
    if Hold_method == "Equal":
        Weights = [(1 / len(Portfolio))]* len(Portfolio)
        # print(Weights)
    elif Hold_method == "Adapt":
        df = Sel_from_daily_ret(Start, Start, Stock_code = Portfolio)
        Sum = df["Dsmvtll"].sum()
        for item in Portfolio:
            temp = (df[df["Stkcd"] == item]["Dsmvtll"].sum()) / Sum
            Weights.append(temp)
        # print(Weights)


    ## 根据持仓周期，获得持仓日历

    price_df = pd.DataFrame()

    if Hold_days != 0:
        Time_use = []
        for i in range(Hold_days + 1):
            Time_use.append(Time_list[i])
        Time_list = Time_use

    elif Hold_period == "Year":
        Start_year = Start_time.year

        Time_use = []
        for i in range(0, len(Time_list)):
            if Time_list[i].year == Start_year:
                Time_use.append(Time_list[i])

        Time_list = Time_use

    elif Hold_period == "Season":
        Start_month =Start_time.month
        End_month = Start_month + 2

        Time_use = []
        for i in range(0, len(Time_list)):
            if Time_list[i].month >= Start_month and Time_list[i].month <= End_month:
                Time_use.append(Time_list[i])
            else:
                break
        Time_list = Time_use

    elif Hold_days == "Month":
        Start_month = Start_time.month

        Time_use = []
        for i in range(0, len(Time_list)):
            if Time_list[i].month == Start_month:
                Time_use.append(Time_list[i])
            else:
                break
        Time_list = Time_use


    RangeL = Time_list[0].strftime("%Y%m%d")
    RangeR = Time_list[-1].strftime("%Y%m%d")

    ##计算每日收益率
    price_df = Sel_from_daily_ret(RangeL, RangeR, Stock_code=Portfolio)
    price_df["Stock_ret"] = price_df["Dretwd"] + 1

    ## 计算第一天的收益率， Close/open

    Cost_df = Sel_from_daily_ret(Start, Start, Stock_code=Portfolio)
    Cost = []
    for item in Portfolio:
        temp = ((Cost_df[Cost_df["Stkcd"] == item]["Clsprc"]) / (Cost_df[Cost_df["Stkcd"] == item]["Opnprc"])).sum()
        Cost.append(temp)

    ## 修正第一日的收益率
    new_price_df = []
    count = 0
    for item in Portfolio:
        temp = price_df[price_df["Stkcd"] == item]
        temp.sort_values("Trddt", inplace = True)
        temp["Stock_ret"][0] = Cost[count]
        count += 1
        new_price_df.append(temp)

    price_df = pd.concat(new_price_df)
    price_df = price_df.reset_index(drop = True)

    ## 根据每日个股收益加权计算组合收益
    ret_df = price_df.groupby("Trddt").apply(lambda x : np.average(x["Stock_ret"], weights=Weights)).to_frame()
    ret_df.columns = ["Daily_ret"]
    ret_df = ret_df.reset_index(drop = False)

    ## 计算累计收益率
    ret_df["Acc_ret"] = ret_df["Daily_ret"].cumprod()
    ret_df["Acc_ret"][0] = ret_df["Daily_ret"][0]
    return ret_df

def Bench_mark(Start:str, Hold_days : int = 0, Hold_period : str = ''):
    '''
        用于回测基准指数， 包括必选参数 Start， 以及可选参数 Mark_index，Hold_days 和 Hold_period
        基准指数为沪深300
        Start：str  回测的开始日期，形式为 YYYYMMDD
        Hold_days : int 以日为单位回测的持股周期，需要大于等于 1，Start 日买入， 持有Hold_days个交易日，收盘价卖出
        Hold_period: str 以自然月"Month"/季度"Season"/年"Year"为单位持股周期，在周期的最后一天卖出，周期的第一天买入
        注意，季度持仓默认不跨年，如果需要跨年，可以自己用month连续做三次。
        只循环一个周期
        返回的 DataFrame 包括两列， Trddt， Daily_ret, Acc_ret
    '''
    HS300 = pd.read_csv("~/Desktop/mpacc2/data_invest/temp/data_lib/hs300.csv", encoding="gbk")
    HS300 = pd.DataFrame(HS300)
    HS300["Trddt"] = pd.to_datetime(HS300["Trddt"])

    calen = Get_trade_dates(Start, '20221231')
    Start_time = pd.to_datetime(Start, format = "%Y%m%d")

    calen = calen[calen["cal_date"] >= Start_time]
    calen = calen.sort_values("cal_date", ascending = True)
    calen = calen.reset_index(drop = True)
    Time_list = list(calen["cal_date"])

    if Hold_days != 0:
        Time_use = []
        for i in range(Hold_days + 1):
            Time_use.append(Time_list[i])
        Time_list = Time_use

    elif Hold_period == "Year":
        Start_year = Start_time.year

        Time_use = []
        for i in range(0, len(Time_list)):
            if Time_list[i].year == Start_year:
                Time_use.append(Time_list[i])

        Time_list = Time_use

    elif Hold_period == "Season":
        Start_month =Start_time.month
        End_month = Start_month + 2

        Time_use = []
        for i in range(0, len(Time_list)):
            if Time_list[i].month >= Start_month and Time_list[i].month <= End_month:
                Time_use.append(Time_list[i])
            else:
                break
        Time_list = Time_use

    elif Hold_period == "Month":
        Start_month = Start_time.month

        Time_use = []
        for i in range(0, len(Time_list)):
            if Time_list[i].month == Start_month:
                Time_use.append(Time_list[i])
            else:
                break
        Time_list = Time_use

    RangeL = Time_list[0]
    RangeR = Time_list[-1]

    HS300.sort_values("Trddt", ascending = True, inplace = True)

    HS300 = HS300[HS300["Trddt"] >= RangeL]
    HS300 = HS300[HS300["Trddt"] <= RangeR]
    HS300.sort_values("Trddt", ascending = True, inplace = True)
    HS300 = HS300.reset_index(drop = True)

    HS300["Daily_ret"] = HS300["Pct_change"] * 0.01 + 1

    HS300["Daily_ret"][0] = HS300["Close"][0] / HS300["Open"][0]
    HS300["Acc_ret"] = HS300["Daily_ret"].cumprod()

    flag_col = ["Trddt", "Daily_ret", "Acc_ret"]
    return HS300[flag_col]


if __name__ == "__main__":
    show_df = Bench_mark("20150101", Hold_period= "Month")
    show_df.to_csv("test1.csv")
